#!/bin/bash

# --- 1. 定义固定目录 ---
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
FIXED_CACHE_DIR="$SCRIPT_DIR/cache"
FIXED_RESULT_DIR="$SCRIPT_DIR/result"

OS_TYPE="$(uname)"

if [[ "$OS_TYPE" == "Darwin" ]]; then
    echo "--- 检测到 macOS ---"
    # 检查是否已安装 Trivy
    if command -v trivy &> /dev/null; then
        echo "--- 检测到系统已安装 Trivy ---"
        read -p "Trivy 已安装，是否卸载后重新安装？(y/n): " reinstall_choice
        if [[ "$reinstall_choice" == "y" || "$reinstall_choice" == "Y" ]]; then
            echo "--- 正在卸载旧版本 Trivy... ---"
            brew list trivy &>/dev/null && brew uninstall trivy
            sudo rm -f /usr/local/bin/trivy
            echo "--- 旧的 Trivy 组件已清理 ---"
        else
            echo "--- 已取消安装，脚本退出 ---"
            exit 0
        fi
    fi
    # 安装 Trivy
    if command -v brew &>/dev/null; then
        echo "--- 使用 Homebrew 安装 Trivy ---"
        brew install aquasecurity/trivy/trivy
    else
        echo "--- 未检测到 Homebrew，尝试下载安装包 ---"
        TRIVY_VERSION="0.51.1"
        TRIVY_PKG="trivy_${TRIVY_VERSION}_macOS-64bit.tar.gz"
        curl -LO "https://github.com/aquasecurity/trivy/releases/download/v${TRIVY_VERSION}/${TRIVY_PKG}"
        tar zxvf "$TRIVY_PKG"
        sudo mv trivy /usr/local/bin/
        rm -f "$TRIVY_PKG"
    fi
    # --- 创建并配置 Trivy 的缓存和结果目录 ---
    echo "--- 正在创建和配置 Trivy 目录... ---"
    sudo mkdir -p "$FIXED_CACHE_DIR"
    sudo mkdir -p "$FIXED_RESULT_DIR"
    # macOS 下 chown 用户:组，组用id -gn
    sudo chown -R $(whoami):$(id -gn) "$FIXED_CACHE_DIR"
    sudo chown -R $(whoami):$(id -gn) "$FIXED_RESULT_DIR"
    echo "--- 目录创建和权限设置完成 ---"
    # --- 设置环境变量 (zsh优先) ---
    echo "--- 正在设置 Trivy 环境变量... ---"
    if [ -f "$HOME/.zshrc" ]; then
        SHELL_RC="$HOME/.zshrc"
    else
        SHELL_RC="$HOME/.bash_profile"
    fi
    grep -q '^export TRIVY_CACHE_DIR=' "$SHELL_RC" && sed -i '' '/^export TRIVY_CACHE_DIR=/d' "$SHELL_RC"
    grep -q '^export TRIVY_DB_DIR=' "$SHELL_RC" && sed -i '' '/^export TRIVY_DB_DIR=/d' "$SHELL_RC"
    echo "export TRIVY_CACHE_DIR=\"$FIXED_CACHE_DIR\"" >> "$SHELL_RC"
    echo "export TRIVY_DB_DIR=\"$FIXED_CACHE_DIR/db\"" >> "$SHELL_RC"
    export TRIVY_CACHE_DIR="$FIXED_CACHE_DIR"
    export TRIVY_DB_DIR="$FIXED_CACHE_DIR/db"
    echo "--- 环境变量设置完成 ---"
    # --- 运行一次示例扫描 ---
    echo "--- 正在运行一次示例 Trivy 扫描... ---"
    TRIVY_BIN="$(command -v trivy 2>/dev/null || echo /usr/local/bin/trivy)"
    if [ -x "$TRIVY_BIN" ]; then
        "$TRIVY_BIN" fs / --format table --output "$FIXED_RESULT_DIR/rootfs_scan_$(date +%Y%m%d_%H%M%S).txt"
    else
        echo "--- 警告：trivy 命令未找到，示例扫描未执行。请重新登录终端或手动执行 trivy。---"
    fi
    echo "--- Trivy 安装和配置脚本执行完毕 ---"
    echo "请检查 $FIXED_RESULT_DIR 目录下的扫描结果文件。"
    echo "下次登录时，Trivy 环境变量将自动生效。"
    exit 0
fi

# --- 检查是否已安装 Trivy ---
if command -v trivy &> /dev/null; then
    echo "--- 检测到系统已安装 Trivy ---"
    read -p "Trivy 已安装，是否卸载后重新安装？(y/n): " reinstall_choice
    if [[ "$reinstall_choice" == "y" || "$reinstall_choice" == "Y" ]]; then
        echo "--- 正在卸载旧版本 Trivy... ---"
        sudo apt-get remove trivy -y
        sudo rm -rf /usr/local/bin/trivy
        sudo rm -f /etc/apt/sources.list.d/trivy.list
        sudo rm -f /usr/share/keyrings/trivy.gpg
        echo "--- 旧的 Trivy 组件已清理 ---"
    else
        echo "--- 已取消安装，脚本退出 ---"
        exit 0
    fi
fi

# --- 检查本地安装包 ---
LOCAL_DEB=$(ls "$SCRIPT_DIR"/trivy*.deb 2>/dev/null | head -n1)
LOCAL_BIN="$SCRIPT_DIR/trivy"
if [ -f "$LOCAL_DEB" ] || [ -f "$LOCAL_BIN" ]; then
    echo "--- 检测到本地 Trivy 安装包 ---"
    echo "1. 使用本地安装包安装 (优先deb包)"
    echo "2. 在线安装 (apt)"
    read -p "请选择安装方式 [1/2]: " install_mode
    if [ "$install_mode" = "1" ]; then
        if [ -f "$LOCAL_DEB" ]; then
            echo "--- 使用本地deb包安装: $LOCAL_DEB ---"
            sudo dpkg -i "$LOCAL_DEB"
            sudo apt-get install -f -y
        elif [ -f "$LOCAL_BIN" ]; then
            echo "--- 使用本地二进制文件安装: $LOCAL_BIN ---"
            sudo install -m 755 "$LOCAL_BIN" /usr/local/bin/trivy
        fi
        echo "--- Trivy 离线安装完成 ---"
    else
        echo "--- 选择在线安装 ---"
        # --- 2. 卸载现有 Trivy (如果存在) ---
        if command -v trivy &> /dev/null; then
            echo "--- 检测到已安装 Trivy，正在卸载旧版本... ---"
            sudo apt-get remove trivy -y
            sudo rm -rf /usr/local/bin/trivy # 删除手动安装的二进制文件
            sudo rm -f /etc/apt/sources.list.d/trivy.list
            sudo rm -f /usr/share/keyrings/trivy.gpg
            echo "--- 旧的 Trivy 组件已清理 ---"
        else
            echo "--- 未检测到已安装 Trivy，跳过卸载步骤 ---"
        fi

        # --- 3. 安装 Trivy ---
        echo "--- 正在安装 Trivy... ---"
        sudo apt-get update
        sudo apt-get install wget apt-transport-https gnupg lsb-release -y

        # 导入 Trivy 的 GPG 密钥
        wget -qO - https://aquasecurity.github.io/trivy-repo/deb/public.key | gpg --dearmor | sudo tee /usr/share/keyrings/trivy.gpg > /dev/null

        # 添加 Trivy 仓库到 APT 源列表
        echo "deb [signed-by=/usr/share/keyrings/trivy.gpg] https://aquasecurity.github.io/trivy-repo/deb $(lsb_release -sc) main" | sudo tee /etc/apt/sources.list.d/trivy.list

        # 更新软件包列表并安装 Trivy
        sudo apt-get update
        sudo apt-get install trivy -y
        echo "--- Trivy 安装完成 ---"
    fi
else
    # --- 2. 卸载现有 Trivy (如果存在) ---
    if command -v trivy &> /dev/null; then
        echo "--- 检测到已安装 Trivy，正在卸载旧版本... ---"
        sudo apt-get remove trivy -y
        sudo rm -rf /usr/local/bin/trivy # 删除手动安装的二进制文件
        sudo rm -f /etc/apt/sources.list.d/trivy.list
        sudo rm -f /usr/share/keyrings/trivy.gpg
        echo "--- 旧的 Trivy 组件已清理 ---"
    else
        echo "--- 未检测到已安装 Trivy，跳过卸载步骤 ---"
    fi

    # --- 3. 安装 Trivy ---
    echo "--- 正在安装 Trivy... ---"
    sudo apt-get update
    sudo apt-get install wget apt-transport-https gnupg lsb-release -y

    # 导入 Trivy 的 GPG 密钥
    wget -qO - https://aquasecurity.github.io/trivy-repo/deb/public.key | gpg --dearmor | sudo tee /usr/share/keyrings/trivy.gpg > /dev/null

    # 添加 Trivy 仓库到 APT 源列表
    echo "deb [signed-by=/usr/share/keyrings/trivy.gpg] https://aquasecurity.github.io/trivy-repo/deb $(lsb_release -sc) main" | sudo tee /etc/apt/sources.list.d/trivy.list

    # 更新软件包列表并安装 Trivy
    sudo apt-get update
    sudo apt-get install trivy -y
    echo "--- Trivy 安装完成 ---"
fi

# --- 4. 创建并配置 Trivy 的缓存和结果目录 ---
echo "--- 正在创建和配置 Trivy 目录... ---"
sudo mkdir -p "$FIXED_CACHE_DIR"
sudo mkdir -p "$FIXED_RESULT_DIR"

# 确保当前用户有写入权限
sudo chown -R $(whoami):$(whoami) "$FIXED_CACHE_DIR"
sudo chown -R $(whoami):$(whoami) "$FIXED_RESULT_DIR"
echo "--- 目录创建和权限设置完成 ---"

# --- 5. 设置环境变量 (永久生效) ---
echo "--- 正在设置 Trivy 环境变量... ---"
# 移除旧的 Trivy 环境变量设置，避免重复
sed -i '/^export TRIVY_CACHE_DIR=/d' ~/.bashrc
sed -i '/^export TRIVY_DB_DIR=/d' ~/.bashrc

echo "export TRIVY_CACHE_DIR=\"$FIXED_CACHE_DIR\"" >> ~/.bashrc
echo "export TRIVY_DB_DIR=\"$FIXED_CACHE_DIR/db\"" >> ~/.bashrc

# 重新加载 .bashrc 以使环境变量立即生效
source ~/.bashrc
echo "--- 环境变量设置完成 ---"

# --- 6. 运行一次示例扫描 ---
echo "--- 正在运行一次示例 Trivy 扫描... ---"
TRIVY_BIN="$(command -v trivy 2>/dev/null || echo /usr/local/bin/trivy)"
if [ -x "$TRIVY_BIN" ]; then
    "$TRIVY_BIN" fs / --format table --output "$FIXED_RESULT_DIR/rootfs_scan_$(date +%Y%m%d_%H%M%S).txt"
else
    echo "--- 警告：trivy 命令未找到，示例扫描未执行。请重新登录终端或手动执行 trivy。---"
fi

echo "--- Trivy 安装和配置脚本执行完毕 ---"
echo "请检查 $FIXED_RESULT_DIR 目录下的扫描结果文件。"
echo "下次登录时，Trivy 环境变量将自动生效。"